// (C) Copyright to Kevin. Applied to MIT LICENSE

#include "NetworkClient.h"

#include "../BlockWorld.h"

#include "NetworkingCommon.h"
#include "TcpSocketBuilder.h"
#include "NetworkRW.h"
#include "GZipCodec.h"
#include "URKProto.h"

NetworkClient::NetworkClient()
{
	InitializeHandlers();
}

NetworkClient::~NetworkClient()
{
	socket->Close();
	thread->Kill(false);
}

void NetworkClient::Start()
{
	thread = FRunnableThread::Create(this, TEXT("BlockEngine Network Client"), 0, EThreadPriority::TPri_AboveNormal);
}

#define REGISTER_HANDLER(id, handler_body) \
  message_handlers.Add(id, handler_body);

void NetworkClient::InitializeHandlers() {
	// Generated by URK_PB_CONVERTER
	REGISTER_HANDLER(MESSAGE_ID_ServerHandshakeMessage, &message_handler_ServerHandshakeMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerAuthenticateResultMessage, &message_handler_ServerAuthenticateResultMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerDisconnectMessage, &message_handler_ServerDisconnectMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerJoinGameMessage, &message_handler_ServerJoinGameMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdateEnvironmentMessage, &message_handler_ServerUpdateEnvironmentMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerChunkMessage, &message_handler_ServerChunkMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdateBlockMessage, &message_handler_ServerUpdateBlockMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdateBlockBatchMessage, &message_handler_ServerUpdateBlockBatchMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerChatMessage, &message_handler_ServerChatMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdatePlayerPositionMessage, &message_handler_ServerUpdatePlayerPositionMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerAddEntityMessage, &message_handler_ServerAddEntityMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerEntityUpdateMessage, &message_handler_ServerEntityUpdateMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerRemoveEntityMessage, &message_handler_ServerRemoveEntityMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerClearEntitiesMessage, &message_handler_ServerClearEntitiesMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerEntityBindingControlMessage, &message_handler_ServerEntityBindingControlMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerWindowOpenMessage, &message_handler_ServerWindowOpenMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerWindowCloseMessage, &message_handler_ServerWindowCloseMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdateWindowMessage, &message_handler_ServerUpdateWindowMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerUpdateWindowElementMessage, &message_handler_ServerUpdateWindowElementMessage);
	REGISTER_HANDLER(MESSAGE_ID_ServerCursorItemMessage, &message_handler_ServerCursorItemMessage);
	// End of generated code
}

uint32 NetworkClient::Run()
{
	UE_LOG(URKNetworking, Log, TEXT("Networking starting... "));
	ISocketSubsystem* subsys = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);
	//socket = subsys->CreateSocket(NAME_DGram, TEXT("UnrealKingdom"), true);
	socket = FTcpSocketBuilder(TEXT("UnrealKingdom")).AsBlocking().Build();

	UE_LOG(URKNetworking, Log, TEXT("Connecting... "));
	TSharedPtr<FInternetAddr> target = subsys->CreateInternetAddr();
	bool valid;
	target->SetIp(TEXT("127.0.0.1"), valid);
	UE_LOG(URKNetworking, Log, TEXT("IP Address validation: %s"), valid ? TEXT("OKAY") : TEXT("ERROR"));
	target->SetPort(25999);
	last_communicate_time = GetCurrentMillisecondTimestamp();
	bool connection_success = socket->Connect(*target);
	if (!connection_success) {
		connection_failed = true;
		UE_LOG(URKNetworking, Log, TEXT("Connection failed! Current state = %d. "), GetConnectionState());
		if (cb_on_connection_failed) cb_on_connection_failed();
		return 0;
	}
	connected_time = GetCurrentMillisecondTimestamp();

	if (cb_on_connected) cb_on_connected();

	UE_LOG(URKNetworking, Log, TEXT("Connection success! Current state = %d. "), GetConnectionState());

	uint8* buffer_header = new uint8[8]; // type and size
	uint8* buffer_body = (uint8*) FMemory::Malloc(sizeof(uint8) * 1024 * 1024 * 16);

	uint32 type = 0;
	uint32 size = 0;

	while (GetConnectionState() == CONNECTION_STATE_CONNECTED) {
		int32 read = -1;
		{
			read_header:
			if (!socket->Wait(ESocketWaitConditions::WaitForRead, FTimespan(100))) {
				goto read_header;
			} else {
				uint32 pending;
				bool has_data = socket->HasPendingData(pending);
				if (!has_data || pending < 8) {
					FPlatformProcess::Sleep(.025f);
					goto read_header;
				}
			}
			bool success = socket->Recv(buffer_header, 8, read, ESocketReceiveFlags::WaitAll);
			if (!success) { // disconnected or error? 
				disconnected = true;
				UE_LOG(URKNetworking, Log, TEXT("Socket error, disconnected! "));
				if (cb_on_disconnected) cb_on_disconnected();
				goto exit_loop;
			}
			type = ReadUInt32(buffer_header, 0);
			size = ReadUInt32(buffer_header, 4);
		}
		UE_LOG(URKNetworking, Log, TEXT("Start reading packet with type 0x%08x and size = %d. "), type, size);
		//if (FMemory::GetAllocSize(buffer_body) < size) {
		//	FMemory::Realloc(buffer_body, size);
		//	UE_LOG(URKNetworking, Log, TEXT("Buffer too small, reallocated. "));
		//}
		if(size > 0) {
			read_body: // LABEL
			if (!socket->Wait(ESocketWaitConditions::WaitForRead, FTimespan(100))) {
				goto read_body;
			} else {
				uint32 pending;
				bool has_data = socket->HasPendingData(pending);
				if (!has_data || (uint32) pending < size) {
					FPlatformProcess::Sleep(.025f);
					goto read_body;
				}
			}
			bool success = socket->Recv(buffer_body, size, read, ESocketReceiveFlags::WaitAll);
			if (!success) { // disconnected or error? 
				disconnected = true;
				UE_LOG(URKNetworking, Log, TEXT("Socket error, disconnected! "));
				if (cb_on_disconnected) cb_on_disconnected();
				goto exit_loop;
			}
		}

		// Packet received
		std::string compressed((const char *)buffer_body, size);
		std::string inflated;
		if (!GZipCodec::Uncompress(compressed, inflated)) {
			UE_LOG(URKNetworking, Log, TEXT("[X] Failed uncompressing message 0x%08x. "), type);
			continue;
		}
		UE_LOG(URKNetworking, Log, TEXT("[#] Uncompressed message 0x%08x body size = %d. "), type, inflated.size());
		google::protobuf::Message* msg = CreateMessageByID(type, (uint8*) inflated.data(), inflated.size()); // new
		if (!msg) {
			UE_LOG(URKNetworking, Log, TEXT("[X] Failed decoding message 0x%08x. "), type);
			continue;
		}
		UE_LOG(URKNetworking, Log, TEXT("[<] Received message type %s. "), UTF8_TO_TCHAR(msg->GetDescriptor()->name().c_str()));
		if (message_handlers.Contains(type)) {
			const MessageHandler** func = message_handlers.Find(type);
			if (*func != nullptr && **func) {
				MessageProcessTask task(*func, msg);
				ABlockEngineGameMode::GetInstance()->GetThreadPool()->AddQueuedWork(&task);
			};
		}
	}

	exit_loop:

	delete[] buffer_header;
	FMemory::Free(buffer_body);

	return 0;
}

int NetworkClient::GetConnectionState()
{
	if (!socket) return CONNECTION_STATE_DISCONNECTED;
	if (connection_failed) return CONNECTION_STATE_FAILED;
	if (disconnected) return CONNECTION_STATE_DISCONNECTED;
	uint64 current = GetCurrentMillisecondTimestamp();
	if (current - last_communicate_time > 10000) return CONNECTION_STATE_TIMED_OUT;
	return CONNECTION_STATE_CONNECTED;
}


void NetworkClient::Shutdown()
{
	if (thread) {
		thread->Kill(false);
		thread = NULL;
	}
	connected_time = 0;
	bool succ = socket->Close();
	UE_LOG(URKNetworking, Log, TEXT("Closing the socket: %s. "), succ ? TEXT("SUCCESS") : TEXT("FAILED"));
	ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM)->DestroySocket(socket);
	socket = NULL;
}

void NetworkClient::SendMessage(const google::protobuf::Message& message)
{
	uint32 id = GetMessageID(message);
	if (id == MESSAGE_ID_UNKNOWN) {
		UE_LOG(URKNetworking, Log, TEXT("Unknown message ID for message %s! "), UTF8_TO_TCHAR(message.GetDescriptor()->name().c_str()));
		return;
	}
	UE_LOG(URKNetworking, Log, TEXT("[>] Sending message %s(ID[%d])"), UTF8_TO_TCHAR(message.GetDescriptor()->name().c_str()), id);
	std::string binary = message.SerializeAsString();
	std::string compressed;
	GZipCodec::Compress(binary, compressed);
	uint8* data = (uint8*) compressed.c_str();
	SIZE_T size = compressed.size();

	uint8* buffer = (uint8*) FMemory::Malloc(8 + size);
	WriteUInt32(buffer, 0, id);
	WriteUInt32(buffer, 4, size);
	FMemory::Memcpy(buffer + 8, data, size);
	int32 sent = 0;
	bool success = socket->Send(buffer, 8 + size, sent);
	if (!success || sent != size + 8) {
		UE_LOG(URKNetworking, Log, TEXT("[X] Send message failed, size = %d, sent = %d. "), size + 8, sent);
	}
}

#define REGISTER_OUT(type_name, id) \
  if (strcmp(message.GetDescriptor()->name().c_str(), type_name) == 0) return (uint32) id;

uint32 NetworkClient::GetMessageID(const google::protobuf::Message& message) {
	// UE_LOG(URKNetworking, Log, TEXT("Getting message ID for %s(%d). "), UTF8_TO_TCHAR(message->GetDescriptor()->name().c_str()), message->GetDescriptor()->name().length());

	// Generated by URK_PB_CONVERTER
	REGISTER_OUT("ClientHandshakeMessage", 0xAABBCCDD);
	REGISTER_OUT("ClientAuthenticateMessage", 0xAF000000);
	REGISTER_OUT("ClientChatMessage", 0xE0000000);
	REGISTER_OUT("ClientMovementMessage", 0xE0000001);
	REGISTER_OUT("ClientActionMessage", 0xE0000002);
	REGISTER_OUT("ClientRemoveBlockMessage", 0xE0000003);
	REGISTER_OUT("ClientUseItemMessage", 0xE0000004);
	REGISTER_OUT("ClientHotbarSelectionMessage", 0xE1000000);
	REGISTER_OUT("ClientPickUpItemMessage", 0xE1000001);
	REGISTER_OUT("ClientWindowInteractMessage", 0xE1FF0000);
	REGISTER_OUT("ClientWindowCloseMessage", 0xE1FEFFFF);
	REGISTER_OUT("ClientEntityInteractMessage", 0xBE000000);
	// End of generated code

	return MESSAGE_ID_UNKNOWN;
}

#define REGISTER_IN(target_id, type_name) \
  if(id == target_id) {\
    urk::pb::type_name* message = new urk::pb::type_name();\
    bool success = message->ParseFromArray(data, size);\
    if(!success) {\
        delete message;\
		return NULL;\
	}\
    return message;\
  }

google::protobuf::Message * NetworkClient::CreateMessageByID(uint32 id, uint8 * data, SIZE_T size)
{
	if (data == NULL) return NULL;
	// Generated by URK_PB_CONVERTER
	REGISTER_IN(0x11223344, ServerHandshakeMessage);
	REGISTER_IN(0xAF000001, ServerAuthenticateResultMessage);
	REGISTER_IN(0xFF000000, ServerDisconnectMessage);
	REGISTER_IN(0xB0000000, ServerJoinGameMessage);
	REGISTER_IN(0xB0000001, ServerUpdateEnvironmentMessage);
	REGISTER_IN(0xB0000002, ServerChunkMessage);
	REGISTER_IN(0xB0000003, ServerUpdateBlockMessage);
	REGISTER_IN(0xB0000004, ServerUpdateBlockBatchMessage);
	REGISTER_IN(0xB0000005, ServerChatMessage);
	REGISTER_IN(0xB0000006, ServerUpdatePlayerPositionMessage);
	REGISTER_IN(0xBE000001, ServerAddEntityMessage);
	REGISTER_IN(0xBE000002, ServerEntityUpdateMessage);
	REGISTER_IN(0xBE000003, ServerRemoveEntityMessage);
	REGISTER_IN(0xBE0000FF, ServerClearEntitiesMessage);
	REGISTER_IN(0xBEA00000, ServerEntityBindingControlMessage);
	REGISTER_IN(0xBA000000, ServerWindowOpenMessage);
	REGISTER_IN(0xBA000001, ServerWindowCloseMessage);
	REGISTER_IN(0xBA000002, ServerUpdateWindowMessage);
	REGISTER_IN(0xBA000003, ServerUpdateWindowElementMessage);
	REGISTER_IN(0xBA000004, ServerCursorItemMessage);
	// End of generated code

	return NULL;
}

MessageProcessTask::MessageProcessTask(const MessageHandler * handler, google::protobuf::Message* message)
	: handler(handler), message(message)
{
}

void MessageProcessTask::DoThreadedWork()
{
	(*handler)(message);
	delete message;
	message = NULL;
}
